#include "vector/one_hot.h"
#include <cstddef>
#include <stdexcept>
#include <algorithm>
#include <string>


std::vector<std::vector<uint8_t>> OneHot::encode(const std::vector<uint8_t> &labels, uint64_t num_classes)
{
    // 创建结果向量，大小与输入标签数量相同
    std::vector<std::vector<uint8_t>> one_hot_labels;
    one_hot_labels.reserve(labels.size());

    for (const auto &label : labels)
    {
        if (label >= num_classes)
        {
            throw std::runtime_error("Invalid label value: " + std::to_string(label));
        }

        // 创建一个全0的向量，大小为类别数
        std::vector<uint8_t> encoded(num_classes, 0);

        // 在对应类别位置设置为1
        encoded[label] = 1;

        // 添加到结果中
        one_hot_labels.push_back(std::move(encoded));
    }

    return one_hot_labels;
}

std::vector<uint8_t> OneHot::decode(const std::vector<std::vector<uint8_t>> &one_hot_labels)
{
    std::vector<uint8_t> labels;
    labels.reserve(one_hot_labels.size());

    for (const auto &one_hot : one_hot_labels)
    {
        // 检查向量是否为空
        if (one_hot.empty())
        {
            throw std::runtime_error("Empty one-hot vector encountered");
        }

        // 找到值为1的元素的索引
        auto max_it = std::max_element(one_hot.begin(), one_hot.end());
        size_t index = std::distance(one_hot.begin(), max_it);

        // 验证是否确实是one-hot编码（只有一个1，其余为0）
        if (*max_it != 1 || std::count(one_hot.begin(), one_hot.end(), 1) != 1)
        {
            throw std::runtime_error("Invalid one-hot encoding: vector should contain exactly one '1'");
        }

        // 检查索引是否在有效范围内（0-9）
        if (index > 9)
        {
            throw std::runtime_error("One-hot index out of range (0-9): " + std::to_string(index));
        }

        labels.push_back(static_cast<uint8_t>(index));
    }

    return labels;
}